{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "8.3_auto-parallel.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YjUdjVJGIO0_",
        "colab_type": "text"
      },
      "source": [
        "# 8.3 自动并行计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hNlHzFeBH8Wk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "# 安装 TensorFlow\n",
        "try:\n",
        "  # Colab only\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "    pass"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EHytkPUkKAQv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import time"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DyNEzpdHILEN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Benchmark(object):\n",
        "  def __init__(self, prefix=None):\n",
        "    self.prefix = prefix + ' ' if prefix else ''\n",
        "\n",
        "  def __enter__(self):\n",
        "    self.start = time.time()\n",
        "\n",
        "  def __exit__(self, *args):\n",
        "    print('%stime: %.4f sec' % (self.prefix, time.time() - self.start))  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ONyMXQyTIucq",
        "colab_type": "code",
        "outputId": "1c524288-57c5-4afe-db59-fa6d5f457837",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Num GPUs Available:  1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iW-iWPtvIYay",
        "colab_type": "text"
      },
      "source": [
        "## 8.3.1 CPU 和 GPU 的并行计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dG1aL1s2IN2M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def run(x):\n",
        "  return [tf.matmul(x, x) for _ in range(10)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kDpNSY0oIk_y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with tf.device('/CPU:0'):\n",
        "  x_cpu = tf.random.uniform(shape=(2000, 2000))\n",
        "  \n",
        "with tf.device('/GPU:0'):\n",
        "  x_gpu = tf.random.uniform(shape=(6000, 6000))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qYS_lM2oJZTM",
        "colab_type": "code",
        "outputId": "8ce6b2af-6fd8-49a5-e95b-5187920541fa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "run(x_cpu)\n",
        "run(x_gpu)\n",
        "\n",
        "with Benchmark('Run on CPU.'):\n",
        "  run(x_cpu)\n",
        "\n",
        "with Benchmark('Then Run on GPU.'):\n",
        "  run(x_gpu)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on CPU. time: 1.2657 sec\n",
            "Then Run on GPU. time: 0.0008 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lPtk-X2SKLNZ",
        "colab_type": "code",
        "outputId": "dd0d34eb-9446-438b-ef6e-7fb6c2159aac",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "with Benchmark('Run on both CPU and GPU in parallel.'):\n",
        "  run(x_cpu)\n",
        "  run(x_gpu)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on both CPU and GPU in parallel. time: 1.2338 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HyDc7dtZN0Pj",
        "colab_type": "text"
      },
      "source": [
        "## 8.3.2 计算和通信的并行计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I2iYA-m_NfTD",
        "colab_type": "code",
        "outputId": "dc4ba00c-216e-43ca-8a03-e8090e3f1ab0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "def copy_to_cpu(x):\n",
        "  with tf.device('/CPU:0'):\n",
        "    return [y for y in x]\n",
        "\n",
        "with Benchmark('Run on GPU.'):\n",
        "  y = run(x_gpu)\n",
        "\n",
        "with Benchmark('Then copy to CPU.'):\n",
        "  copy_to_cpu(y)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on GPU. time: 0.0009 sec\n",
            "Then copy to CPU. time: 0.0001 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OaQMC2YWN_4J",
        "colab_type": "code",
        "outputId": "f74efaa0-424b-4cc9-aa6d-c4fbc5e64391",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "with Benchmark('Run and copy in parallel.'):\n",
        "    y = run(x_gpu)\n",
        "    copy_to_cpu(y)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run and copy in parallel. time: 0.0365 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4MjP8mLW-4H",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}